#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Word文档解析器模块
专门解析Word文档，匹配需求文档。
"""
import json
import os
import re
from typing import Dict, List, Any, Tuple
from docx import Document
from docx.enum.style import WD_STYLE_TYPE
from docx.shared import Pt
from collections import Counter
import time
import logging

logging.basicConfig(level=logging.INFO)

class DocParser:
    """Word文档解析器类"""

    def __init__(self):
        self.heading_styles = {
            'Heading 1': 1,
            'Heading 2': 2,
            'Heading 3': 3,
            'Heading 4': 4,
            'Heading 5': 5,
            'Heading 6': 6,
            'Title': 0,  # 文档标题
        }

        # 中文标题识别模式
        self.chinese_title_patterns = [
            r'^[一二三四五六七八九十]+[、．\.]\s*.*',  # 一、一．二．等
            r'^\d+[、．\.]\s*.*',  # 1、1．2．等
            r'^\(\d+\)\s*$',  # (1)(2)等，纯数字序号
            r'^\d+\.\d+[、．\.]*\s*.*',  # 1.1 2.1等
        ]

        # 图注识别模式（这些不应该被当作标题）
        self.figure_caption_patterns = [
            r'^图\s*\d+[-\d]*.*',  # 图 1-1、图1、图 2-1 等
            r'^图\s*[一二三四五六七八九十\d]+.*',  # 图一、图二等
        ]

        # 标题关键词
        self.title_keywords = [
            '章', '节', '部分', '步骤', '说明', '介绍', '概述', '内容',
            '功能', '操作', '设置', '参数', '校正', '调试', '验证',
            '保存', '使用', '注意', '警告', '提示'
        ]

    def parse_document(self, file_path: str) -> Dict[str, Any]:
        start_time = time.time()
        doc = Document(file_path)
        title = os.path.splitext(os.path.basename(file_path))[0]

        # 改进基准字号计算：获取所有字号的分布
        font_sizes = []
        for para in doc.paragraphs:
            for run in para.runs:
                if run.font.size:
                    font_sizes.append(run.font.size.pt)

        # 计算基准字号：取中位数作为基准，避免极端值影响
        if font_sizes:
            sorted_sizes = sorted(font_sizes)
            n = len(sorted_sizes)
            if n % 2 == 1:
                base_size = sorted_sizes[n // 2]
            else:
                base_size = (sorted_sizes[n // 2 - 1] + sorted_sizes[n // 2]) / 2
        else:
            base_size = 12

        # 解析文档结构
        sections = self._parse_sections(doc, base_size)

        processing_time = time.time() - start_time
        page_count = len(doc.sections)  # 近似

        return {
            "title": title,
            "sections": sections,
            "metadata": {
                "source_filename": os.path.basename(file_path),
                "page_count": page_count,
                "processing_time_seconds": processing_time
            }
        }

    def _is_title(self, para, style_name, base_size) -> Tuple[bool, int]:
        """
        判断段落是否为标题，并返回标题层级
        """
        text = para.text.strip()

        # 首先检查是否为图注，如果是则不是标题
        for pattern in self.figure_caption_patterns:
            if re.match(pattern, text):
                return False, 0

        # 检查是否为注释，如果是则不是标题
        if re.match(r'^\([^)]*注[^)]*\)', text):
            return False, 0

        # 新增规则：检查中文标号+内容形式（如"一、找到校正程序"），直接识别为标题
        if re.match(r'^[一二三四五六七八九十]+[、．\.]\s*.*', text):
            return True, 1

        # 检查是否有标号（优先于字号检查）
        level = 0
        
        # 一级标题：第一章、一、1 引言、1. 引言、1.引言
        if re.match(r'^第[一二三四五六七八九十\d]+章', text) or re.match(r'^[一二三四五六七八九十]+[、．\.]', text) or \
           re.match(r'^\d+\s+[^\d\s。，；：！？]{1,15}$', text) or re.match(r'^\d+\.\s+[^\d\s。，；：！？]{1,15}$', text) or re.match(r'^\d+\.\s*[^\d\s。，；：！？]{1,15}$', text):  # 匹配"1 引言"、"1. 引言"和"1.引言"这样的格式，限制标题长度和标点符号
            level = 1
        # 多级数字标题：1.1、1.1.1、1.1.1.1等（直接识别为标题）
        elif re.match(r'^\d+\.\d+', text) and re.match(r'^\d+\.\d+.*[^\d\s。，；：！？]{1,15}$', text):  # 匹配"1.1"及更复杂的格式，限制标题长度和标点符号
            if re.match(r'^\d+\.\d+\.\d+\.\d+\.\d+\.\d+', text):  # 1.1.1.1.1.1
                level = 6
            elif re.match(r'^\d+\.\d+\.\d+\.\d+\.\d+', text):  # 1.1.1.1.1
                level = 5
            elif re.match(r'^\d+\.\d+\.\d+\.\d+', text):  # 1.1.1.1
                level = 4
            elif re.match(r'^\d+\.\d+\.\d+', text):  # 1.1.1
                level = 3
            else:  # 1.1
                level = 2
        # 单个数字+点号：1.、2.等（需要通过字号判断）
        elif re.match(r'^\d+[、．\.]\s*$', text):  # 匹配"1."、"2."等格式
            # 需要通过字号判断是否为标题
            max_size = 0
            for run in para.runs:
                if run.font.size and run.font.size.pt > max_size:
                    max_size = run.font.size.pt
            size_ratio = max_size / base_size if base_size > 0 else 1
            if size_ratio > 1.0:
                level = 1
            else:
                return False, 0
        
        # 如果有标号，直接返回
        if level > 0:
            return True, level
        
        # 无标号：需要通过字号判断
        max_size = 0
        for run in para.runs:
            if run.font.size and run.font.size.pt > max_size:
                max_size = run.font.size.pt
        size_ratio = max_size / base_size if base_size > 0 else 1
        
        if size_ratio > 1.3:
            return True, 0  # 返回0表示需要在_parse_sections中动态确定层级
        else:
            return False, 0

    def _is_figure_or_table_caption(self, text) -> bool:
        """
        检查是否为图注或表格标题
        """
        figure_caption_patterns = [
            r'^图\s*\d+[-\d]*.*',  # 图 1-1、图1、图 2-1 等
            r'^图\s*[一二三四五六七八九十\d]+.*',  # 图一、图二等
        ]

        table_caption_patterns = [
            r'^表\s*\d+[-\d]*.*',  # 表 1-1、表1、表 2-1 等
            r'^表\s*[一二三四五六七八九十\d]+.*',  # 表一、表二等
        ]

        is_figure = any(re.match(pattern, text) for pattern in figure_caption_patterns)
        is_table_caption = any(re.match(pattern, text) for pattern in table_caption_patterns)
        
        return is_figure or is_table_caption

    def _is_list_item(self, text) -> bool:
        """
        检查是否为列表项
        """
        return bool(re.match(r'^[-*]\s|^\d+\.\s', text))

    def _should_end_paragraph(self, text, current_paragraphs) -> bool:
        """
        检查是否应该结束当前段落
        """
        return (
            text.endswith(('。', '.', '？', '！', '?', '!', ':', '：')) or 
            len(text) > 100 or  # 长文本自动分段
            (len(current_paragraphs) > 0 and len(" ".join(current_paragraphs + [text])) > 200)
        )

    def _determine_title_level(self, level, last_numbered_level) -> int:
        """
        确定标题层级
        """
        if level > 0:  # 有标号标题
            return level
        
        # 无标号标题，根据上一个有标号标题的层级确定
        if last_numbered_level > 0:
            level = last_numbered_level + 1
        else:
            level = 2  # 默认2级
        
        # 限制最大层级为6
        return min(level, 6)

    def _add_current_content_to_section(self, stack, current_paragraphs, current_list_items):
        """
        将当前段落和列表项添加到当前section
        """
        if current_paragraphs:
            if stack:
                stack[-1]["paragraph"].extend(current_paragraphs)
        
        if current_list_items and stack:
            # 将列表项作为段落处理
            stack[-1]["paragraph"].extend(current_list_items)

    def _create_new_section(self, text, level):
        """
        创建新的section
        """
        return {
            "title": text,
            "level": level,
            "paragraph": [],
            "sections": []
        }

    def _update_section_stack(self, stack, sections, new_section):
        """
        更新section栈
        """
        # 调整栈
        while stack and stack[-1]["level"] >= new_section["level"]:
            stack.pop()
        
        if stack:
            stack[-1]["sections"].append(new_section)
        else:
            sections.append(new_section)
        
        stack.append(new_section)

    def _parse_sections(self, doc, base_size) -> List[Dict[str, Any]]:
        """
        解析文档的章节结构
        """
        sections = []
        stack = []
        current_paragraphs = []
        current_list_items = []
        last_numbered_level = 0  # 记录上一个有标号标题的层级

        for para in doc.paragraphs:
            text = para.text.strip()
            if not text: continue

            # 首先检查是否为图注或表格标题
            if self._is_figure_or_table_caption(text):
                if stack:
                    stack[-1]["paragraph"].append(text)
                continue

            # 检查是否为注释
            if re.match(r'^\([^)]*注[^)]*\)', text) or re.match(r'^注\s*[:：]', text):
                if stack:
                    stack[-1]["paragraph"].append(text)
                continue

            # 检查是否为标题
            is_title, level = self._is_title(para, para.style.name, base_size)

            if is_title:
                # 添加当前段落和列表项到当前section
                self._add_current_content_to_section(stack, current_paragraphs, current_list_items)
                current_paragraphs = []
                current_list_items = []

                # 更新上一个有标号标题的层级
                if level > 0:  # 有标号标题
                    last_numbered_level = level
                
                # 确定标题层级
                level = self._determine_title_level(level, last_numbered_level)

                # 创建新section并更新栈
                new_section = self._create_new_section(text, level)
                self._update_section_stack(stack, sections, new_section)
            else:
                # 处理非标题文本
                if para.style.name.startswith('List') or self._is_list_item(text):
                    current_list_items.append(text)
                else:
                    if current_list_items:
                        # 将列表项作为段落
                        current_paragraphs.extend(current_list_items)
                        current_list_items = []
                    current_paragraphs.append(text)

        # 处理表格（作为段落文本处理）
        for table in doc.tables:
            rows = [[cell.text for cell in row.cells] for row in table.rows]
            # 将表格转换为文本表示
            table_text = "\n".join(["\t".join(row) for row in rows])
            if stack:
                stack[-1]["paragraph"].append(f"[表格]\n{table_text}")
            else:
                if not sections:
                    sections.append({
                        "title": "内容",
                        "level": 1,
                        "paragraph": [f"[表格]\n{table_text}"],
                        "sections": []
                    })
                else:
                    sections[-1]["paragraph"].append(f"[表格]\n{table_text}")

        # 清理剩余
        self._add_current_content_to_section(stack, current_paragraphs, current_list_items)

        return sections

    def save_to_json(self, data: Dict[str, Any], output_path: str):
        """将解析结果保存为JSON文件，确保使用UTF-8编码"""
        try:
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            print(f"JSON文件已保存到: {output_path}")
        except Exception as e:
            print(f"保存JSON文件失败: {e}")

    def print_structure(self, data: Dict[str, Any], indent: int = 0):
        """打印文档结构"""
        prefix = "  " * indent
        print(f"{prefix}[文档] {data['title']}")
        def print_sections(sections: List[Dict], level_indent: int):
            for section in sections:
                section_prefix = "  " * level_indent
                print(f"{section_prefix}{'#' * section['level']} {section['title']}")
                for para in section.get('paragraph', []):
                    if len(para) > 50:
                        print(f"{section_prefix}  [段落] {para[:47]}...")
                    else:
                        print(f"{section_prefix}  [段落] {para}")
                if 'sections' in section:
                    print_sections(section['sections'], level_indent + 1)
        print_sections(data['sections'], indent + 1)
